import sys
import os
import cv2
import re

sys.path.append(os.path.expanduser('~/Codes/PaddleOCR'))
from paddleocr import PaddleOCR

def im_show(image):
    plt.imshow(image)
    plt.show()
    
def show_txt(image, det_boxes_scores):
    bifen_txt = []
    for box, txt_score in det_boxes_scores[0]:
        x1, y1 = box[0]
        x1, y1 = int(x1), int(y1)
        x2, y2 = box[2]
        x2, y2 = int(x2), int(y2)
        color = (0, 255, 0)
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        txt, score = txt_score
        txt = txt.replace(' ', '')
        match_jie = re.match(r"第\d+节", txt)
        match_obj = re.match(r"\d+:\d+", txt)
        if match_obj or match_jie:
            bifen_txt.append(txt)
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(image, txt,(x1-10, y1-10),font, 2, (255,0, 0), 3)
    im_show(image)
    return bifen_txt

def detect_txt(image, det_boxes_scores):
    bifen_txt = []
    for box, txt_score in det_boxes_scores[0]:
        x1, y1 = box[0]
        x1, y1 = int(x1), int(y1)
        x2, y2 = box[2]
        x2, y2 = int(x2), int(y2)
        txt, score = txt_score
        txt = txt.replace(' ', '')
        match_jie = re.match(r"第\d+节", txt)
        match_obj = re.match(r"\d+:\d+", txt)
        if match_obj or match_jie:
            bifen_txt.append(txt)
    if(len(bifen_txt) < 3):
        print("missing score")
        while(len(bifen_txt) < 3):
            bifen_txt.append("miss")
    return bifen_txt


ocr_engine = PaddleOCR()

#image_path = '/mnt/data/huangying/datas/比分大屏/现场比分大屏/上半场/live1.ts_imgs/03601.jpeg'
#image = cv2.imread(image_path)
vid_capture = cv2.VideoCapture('rtsp://192.168.2.162/live_stream')

result_txt = open("result2.txt", 'w')
frame_id = 0
sucess, frame = vid_capture.read()
recog_fps = 25
while sucess:
    frame_id += 1
    sucess, frame = vid_capture.read()
    if frame_id % recog_fps == 0:
        det_boxes_scores = ocr_engine.ocr(frame)
        #bifen_txt = show_txt(frame, det_boxes_score)
        bifen_txt = detect_txt(frame, det_boxes_scores)

        result_txt.write("frame:{}\n".format(frame_id))
        for bifen in bifen_txt:
            print(bifen)
            result_txt.write("{}\n".format(bifen))
            result_txt.flush()
result_txt.close()

#boxes, txt_scores, _ = det_boxes_score

